#############################################################################################
#############################################################################################
#
#   Open Source License/Disclaimer, Forecast Systems Laboratory NOAA/OAR/GSD,
#   325 Broadway Boulder, CO 80305
#
#   This software is distributed under the Open Source Definition, which may be
#   found at http://www.opensource.org/.
#   In particular, redistribution and use in source and binary forms, with or
#   without modification, are permitted provided that the following conditions are met:
#
#   - Redistributions of source code must retain this notice, this list of
#     conditions and the following disclaimer.
#
#   - Redistributions in binary form must provide access to this notice, this
#     list of conditions and the following disclaimer, and the underlying source code.
#
#   - All modifications to this software must be clearly documented, and are
#     solely the responsibility of the agent making the modifications.
#
#   - If significant modifications or enhancements are made to this software,
#     the GSD Software Policy Manager (softwaremgr.fsl@noaa.gov) should be notified.
#
#   THIS SOFTWARE AND ITS DOCUMENTATION ARE IN THE PUBLIC DOMAIN AND ARE
#   FURNISHED "AS IS." THE AUTHORS, THE UNITED STATES GOVERNMENT, ITS INSTRUMENTALITIES,
#   OFFICERS, EMPLOYEES, AND AGENTS MAKE NO WARRANTY, EXPRESS OR IMPLIED, AS TO
#   THE USEFULNESS OF THE SOFTWARE AND DOCUMENTATION FOR ANY PURPOSE. THEY ASSUME
#   NO RESPONSIBILITY (1) FOR THE USE OF THE SOFTWARE AND DOCUMENTATION; OR (2) TO PROVIDE
#   TECHNICAL SUPPORT TO USERS.
#############################################################################################
#############################################################################################


import inspect
import numpy as np

class ModelClass(object):
    
    def __init__(self,start_time,stop_time,timestep,history_interval,nx=1,ny=1,nz=1,dx=1,dy=1,offset=0,time_diff_choice='leapfrog'): 
    # The class is called by initializing the model, which includes initializing 
    # initial condition (boundary conditions), model parameters     
        self.Name = 'Model Name'
        self.nx = nx
        self.ny = ny
        self.nz = nz
        self.dx = dx
        self.dy = dy
        self.offset = offset             # additional grid points for spacial differential scheme 
        self.start_time = start_time
        self.stop_time = stop_time
        self.timestep = timestep
        self.time_diff_choice = time_diff_choice
        self.SimTime = np.arange(self.start_time,self.stop_time,self.timestep)
        self.history_start = start_time
        self.history_end = stop_time
        self.history_interval =history_interval
        self.model_output_history()

    # Construct a diagonal error matrix
    def error_matrix ( self, error ):
        # The class is called by constructing a diagonal error matrix     
        raise NotImplementedError(inspect.stack()[0][3])
    
    # Add normal noise on IC
    def model_ic_add_noise(self):
    # The class is called by adding random noise on the IC     
        raise NotImplementedError(inspect.stack()[0][3])
    
    # Initilize pertubation
    def model_init_tl(self):
    # The class is called by initializing model pertubations     
        raise NotImplementedError(inspect.stack()[0][3])
    
    # Initilize gradient
    def model_init_ad(self):
    # The class is called by initializing model gradients     
        raise NotImplementedError(inspect.stack()[0][3])
                
    # Defines the model tendency
    def model_tendency(self) :
    # Defines the model tendency calculation. 
        raise NotImplementedError(inspect.stack()[0][3])
        
    # Defines the model tendency --- Tangent Linear
    def model_tendency_tl(self) :
        # Defines the model perturbation tendency calculation
        raise NotImplementedError(inspect.stack()[0][3])
    
    # Defines the model tendency  --- Adjoint
    def model_tendency_ad(self) :
    # Defines the transpose of the model perturbation tendency calculation. 
        raise NotImplementedError(inspect.stack()[0][3])
        
    # Defines the model output history
    def model_output_history(self) :
        if self.history_interval <= self.timestep :
            self.history_interal = self.timestep
        self.history = np.linspace( self.history_start, self.history_end, int((self.history_end-self.history_start)/self.history_interval+1) )
        
    ###############################################################################
    # leapfrog-algorithm used in several of the above simulations.
    #
    # Integrate an ODE specified by 
    # y' = f(y,t,args), y(t0) = y0, t0 < t < T
    # Using fourth order Runge-Kutta on a specified time grid.

    # f() must return a numpy column array of outputs of the same
    # dimension as that of y. It must take a numpy column array of
    # variable values y, a single float time t, and any other arguments.

    # The output is a numpy array of y values at the times in t of
    # dimensions (len(t),len(y0)).
    ############################################################################### 
    def leapfrog(self):
        trajectory = np.zeros( (len(self.history),)+ self.ic.shape ) 
        x = self.ic.copy()
        trajectory[0,:] = x.copy()
        counter = 1
        
        for i in range(len(self.SimTime)):
            delta_t = self.timestep
            updates = delta_t * self.model_tendency(x, self.SimTime[i])
            x = x + updates
            # Write into trajectory            
            if (self.SimTime[i]+self.timestep) in self.history:
                trajectory[counter,:] = x.copy() 
                counter += 1        

        # Update fields to current states
        self.u_field.data[:], self.v_field.data[:], self.h_field.data[:] = np.split(x,self.VarDim)
        
        return trajectory

    
    # Model Time differential forward -- Tangent Linear
    def leapfrog_tl(self):
        
        trajectory = np.zeros( (len(self.history),)+ self.ic.shape ) 
        trajectory_tl = np.zeros( (len(self.history),)+ self.ic_tl.shape ) 
        x = self.ic.copy()
        x_tl = self.ic_tl.copy()
        trajectory[0,:] = x.copy()
        trajectory_tl[0,:] = x_tl.copy()
        counter = 1

        for i in range(len(self.SimTime)):
            delta_t = self.timestep
            updates, updates_tl = self.model_tendency_tl(x, x_tl, self.SimTime[i])
            x = x + updates * delta_t
            x_tl = x_tl + updates_tl * delta_t
            
            if (self.SimTime[i]+self.timestep) in self.history:
                trajectory[counter,:] = x.copy()
                trajectory_tl[counter,:] = x_tl.copy()
                counter += 1

        return trajectory, trajectory_tl
    
    # Model Time differential backword -- Adjoint

    def leapfrog_ad(self):

        trajectory = np.zeros( (len(self.SimTime)+1,)+ self.ic.shape ) 
        trajectory[0,:] = self.ic.copy()
        
        for i in range(len(self.SimTime)):
            delta_t = self.timestep
            updates = delta_t * self.model_tendency(trajectory[i,:], self.SimTime[i])
            trajectory[i+1,:] = trajectory[i,:] + updates        
  
        x_ad = np.zeros_like( self.ic_ad )
        updates_ad = np.zeros_like( self.ic_ad )
        counter = len(self.history) - 1

        for i in reversed(range(len(self.SimTime))):
            delta_t = self.timestep 

            if (self.SimTime[i]+self.timestep) in self.history:
                x_ad = x_ad + self.forcing_ad[counter,:].reshape(self.ic_ad.shape)
                self.forcing_ad[counter,:].fill(0.)
                counter -= 1
            updates_ad = x_ad.copy()
            result1b = delta_t * updates_ad
            x_ad = x_ad + self.model_tendency_ad(trajectory[i,:], result1b, self.SimTime[i])
        
        x_ad = x_ad + self.forcing_ad[counter,:].reshape(self.ic_ad.shape)
        self.forcing_ad[counter,:].fill(0.)
        self.ic_ad = x_ad.copy() 

                   
    ###############################################################################
    # rk4-algorithm used in several of the above simulations.
    #
    # Integrate an ODE specified by 
    # y' = f(y,t,args), y(t0) = y0, t0 < t < T
    # Using fourth order Runge-Kutta on a specified time grid.
    ###############################################################################
    def rk4(self):
        trajectory = np.zeros( (len(self.history),)+ self.ic.shape ) 
        x = self.ic.copy()
        trajectory[0,:] = x.copy() 
        counter = 1        
        
        for i in range(len(self.SimTime)):
            delta_t = self.timestep
            k1 = delta_t * self.model_tendency(x, self.SimTime[i])
            k2 = delta_t * self.model_tendency(x + .5*k1, self.SimTime[i] + .5*delta_t)
            k3 = delta_t * self.model_tendency(x + .5*k2, self.SimTime[i] + .5*delta_t)
            k4 = delta_t * self.model_tendency(x + k3, self.SimTime[i] + delta_t)
            x = x + 0.16666666*(k1 + 2*k2 + 2*k3 + k4)
            
            if (self.SimTime[i]+self.timestep) in self.history:
                trajectory[counter,:] = x.copy() 
                counter += 1        
        
        # Update fields to current states
        self.u_field.data[:], self.v_field.data[:], self.h_field.data[:] = np.split(x,self.VarDim)
                
        return trajectory

    def rk4_tl(self):
        # Model Time differential rk4 -- Tangent Linear
        trajectory = np.zeros( (len(self.history),)+ self.ic.shape ) 
        trajectory_tl = np.zeros( (len(self.history),)+ self.ic_tl.shape ) 
        x = self.ic.copy()
        x_tl = self.ic_tl.copy()
        trajectory[0,:] = x.copy()
        trajectory_tl[0,:] = x_tl.copy()
        counter = 1
        
        for i in range(len(self.SimTime)):
            delta_t = self.timestep
            
            k1,k1_tl = self.model_tendency_tl(x, x_tl, self.SimTime[i])
            k1 = delta_t * k1
            k1_tl = delta_t * k1_tl
            k2,k2_tl = self.model_tendency_tl(x + .5*k1, x_tl + .5*k1_tl, self.SimTime[i] + .5*delta_t)
            k2 = delta_t * k2
            k2_tl = delta_t * k2_tl
            k3,k3_tl = self.model_tendency_tl(x + .5*k2, x_tl + .5*k2_tl, self.SimTime[i] + .5*delta_t)
            k3 = delta_t * k3
            k3_tl = delta_t * k3_tl
            k4,k4_tl = self.model_tendency_tl(x + k3, x_tl + k3_tl, self.SimTime[i] + delta_t)
            k4 = delta_t * k4
            k4_tl = delta_t * k4_tl
            x = x + 0.16666666*(k1 + 2*k2 + 2*k3 + k4)
            x_tl = x_tl + 0.16666666*(k1_tl + 2*k2_tl + 2*k3_tl + k4_tl)
            
            if (self.SimTime[i]+self.timestep) in self.history:
                trajectory[counter,:] = x.copy()
                trajectory_tl[counter,:] = x_tl.copy()
                counter += 1

        return trajectory, trajectory_tl
        
    def rk4_ad(self):
        # Model Time differential rk4 -- Adjoint            
        trajectory = np.zeros( (len(self.SimTime)+1,)+ self.ic.shape ) 
        k1_save = np.zeros( (len(self.SimTime),)+ self.ic.shape )
        k2_save = np.zeros( (len(self.SimTime),)+ self.ic.shape )
        k3_save = np.zeros( (len(self.SimTime),)+ self.ic.shape )
        trajectory[0,:] = self.ic.copy()
        
        for i in range(len(self.SimTime)):
            delta_t = self.timestep

            k1 = delta_t * self.model_tendency(trajectory[i,:], self.SimTime[i])
            k1_save[i,:] = k1.copy()
            k2 = delta_t * self.model_tendency(trajectory[i,:] + .5*k1, self.SimTime[i] + .5*delta_t)
            k2_save[i,:] = k2.copy()
            k3 = delta_t * self.model_tendency(trajectory[i,:] + .5*k2, self.SimTime[i] + .5*delta_t)
            k3_save[i,:] = k3.copy()
            k4 = delta_t * self.model_tendency(trajectory[i,:] + k3, self.SimTime[i] + delta_t)
            trajectory[i+1,:] = trajectory[i,:] + 0.16666666*(k1 + 2*k2 + 2*k3 + k4)
            
        x_ad = np.zeros( self.ic_ad.shape )
        arg10b = np.zeros_like( self.ic_ad )
        arg1b = np.zeros_like( self.ic_ad )
        counter = len(self.history) - 1
        
        for i in reversed(range(len(self.SimTime))):
            delta_t = self.timestep 
            
            if (self.SimTime[i]+self.timestep) in self.history:
                x_ad = x_ad + self.forcing_ad[counter,:].reshape(self.ic_ad.shape)
                self.forcing_ad[counter,:].fill(0.)
                counter -= 1
            tempb = 0.16666666 * x_ad
            k4_ad = tempb  
            result1b = delta_t * k4_ad
            arg10b.fill(0.)
            arg10b = self.model_tendency_ad(trajectory[i,:] + k3_save[i,:], result1b, self.SimTime[i] + delta_t)
            k3_ad = arg10b + 2*tempb
            result1b = delta_t * k3_ad
            arg1b.fill(0.)
            arg1b = self.model_tendency_ad(trajectory[i,:] + 0.5*k2_save[i,:], result1b, self.SimTime[i] + 0.5*delta_t)
            k2_ad = 0.5*arg1b + 2*tempb
            x_ad = x_ad + arg1b + arg10b
            result1b = delta_t * k2_ad
            arg1b.fill(0.)
            arg1b = self.model_tendency_ad(trajectory[i,:] + 0.5*k1_save[i,:], result1b, self.SimTime[i] + 0.5*delta_t)
            k1_ad = 0.5*arg1b +tempb
            x_ad = x_ad + arg1b
            result1b = delta_t * k1_ad
            x_ad = x_ad + self.model_tendency_ad(trajectory[i,:], result1b, self.SimTime[i] + delta_t)
              
        x_ad = x_ad + self.forcing_ad[counter,:].reshape(self.ic_ad.shape)
        self.forcing_ad[counter,:].fill(0.)
        self.ic_ad = x_ad.copy() 
               
    def model_integration (self):
        # Model integration
        if self.time_diff_choice.lower() == 'leapfrog':
            Xsim = self.leapfrog()
        elif self.time_diff_choice.lower() == 'rk4':
            Xsim = self.rk4()
        else:
            raise Exception('The time differential choice is either leapfrog or rk4 at this moment!')
        
        return Xsim

    def model_integration_tl (self):
        # Model integration
        if self.time_diff_choice.lower() == 'leapfrog':
            Xsim,Xsim_tl = self.leapfrog_tl()
        elif self.time_diff_choice.lower() == 'rk4':
            Xsim,Xsim_tl = self.rk4_tl()
        else:
            raise Exception('The time differential choice is either leapfrog or rk4 at this moment!')
        
        return Xsim, Xsim_tl

    def model_integration_ad (self):
        # Model integration
        if self.time_diff_choice.lower() == 'leapfrog':
            self.leapfrog_ad()
        elif self.time_diff_choice.lower() == 'rk4':
            self.rk4_ad()
        else:
            raise Exception('The time differential choice is either leapfrog or rk4 at this moment!')
        
        # The final gradient is in self.ic_ad 
                                                
    #  Tangent Linear Check
    def model_tl_check (self):
    
        # Assume we already have a IC and we use ic_add_noise to add 
        # the radom normal noise on the original IC
        self.model_ic_add_noise()
        
        # Initializing model pertubation
        # This is not necessary actually, put it here for clear
        self.model_init_tl()
        
        # Calculate the pertubation
        self.ic_tl = self.ic_perturbation - self.ic
        
        # Integration TLM       
        B_states,G_states = self.model_integration_tl()   
    
        # Calculate the norm of perturbations at timestep = nts      
        dot_tl  = np.sum( np.multiply( G_states[-1,:],G_states[-1:] ) )

        # Save the basic states at timstep = nts      
        End_states = B_states[-1,:].copy()
        
        # Save the IC as it will be used in every iteration
        ic_save = self.ic.copy()
        
        # Call NLM 
        print '==========Tangent Linear Check================'
        alpha_m = 1.
        for order in range(1,12):
            
            alpha_m = 0.1 * alpha_m
            
            # Perturb the initial condition           
            self.ic = ic_save + alpha_m * self.ic_tl
          
            # Call NLM
            B_states = self.model_integration () 

            # Calculate the diffs
            diff = B_states[-1,:] - End_states[:]
            
            # The norm of the basic states at timestep nts  
            dot_nl_new  = np.sum( np.multiply( diff,diff ) )
                  
            print 'tl_check: alpha_m= %9.4e  coef= %18.15f val_n= %14.7e val_l= %14.7e' % \
                ( alpha_m ,dot_nl_new/(alpha_m**2*dot_tl) ,dot_nl_new, alpha_m**2*dot_tl)
    
                
    #  Adjoint Check
    def model_ad_check(self):

        # Initializing model pertubation
        # This is not necessary actually, put it here for clear
        self.model_init_tl()
        
        # Initializing model gradient
        # This is necessary
        self.model_init_ad()
         
        # Assume we already have a IC and we use ic_add_noise to add 
        # the radom normal noise on the original IC
        self.model_ic_add_noise()
               
        # Calculate the pertubation
        self.ic_tl = self.ic_perturbation - self.ic
        
        # Redefine the history output interval
        # We need to save the basic states at every timestep
        self.history_interval=self.timestep
        self.model_output_history()
        
        # Call TLM  
        B_states, G_states = self.model_integration_tl()
        
        # Calclulate the norm of perturbation at timestep = nts
        dot_tl  = np.sum( np.multiply( G_states[-1],G_states[-1] ) )       
        # Take the output of TLM at timestep = nts as the input for ADM       
                
        # Allocate and set adjoint forcing to ZERO
        self.forcing_ad = np.zeros( (len(self.history),)+ self.ic_ad.shape )
        
         # Take the output of TLM at timestep = nts as the input for ADM 
        self.forcing_ad[-1,:] = G_states[-1].copy()
        
        # Call ADM, the  gradient at 0 is in self.ic_ad
        self.model_integration_ad()
    
        # Calclulate the norm of perturbation at timestep = nts
        
        dot_ad  = np.sum( np.multiply( self.ic_tl,self.ic_ad ) )
        
        print '==========  Adjoint Check  ================'
        print 'ad_check: VAL_TL: %23.15f' % dot_tl
        print 'ad_check: VAL_AD: %23.15f' % dot_ad
            
    
